#include "stdafx.h"

namespace classifier_test
{

namespace high_drift
{

modularity_clustering::modularity_clustering()
{
}

void modularity_clustering::clear()
{
	Size = 0;
	Graph.clear();
	reset_vector(Weights);
	NodeCount = 0;
	reset_vector(Nodes);
	Merges.clear();
	Heap.clear();
	ClusCount = 0;
	reset_vector(Labels);
}

void modularity_clustering::set_size(int SetSize)
{
	clear();
	Size = SetSize;
	Graph.set_size(Size);
	Weights.assign(Size, 0);
}

void modularity_clustering::set_division_threshold(double Val)
{
	DivThreshold = Val;
}

void modularity_clustering::add_edge(int A, int B, double Weight)
{
	assert(Weight >= 0);
	if (Weight == 0) return;
	if (A == B) Weights[A] = (weight_type)Weight;
	else Graph.add_edge(A, B, (weight_type)Weight);
}

void modularity_clustering::compute(bool Subdiv)
{
	logger::task LTask = clogger()->fine("Modularity clustering");
	clogger()->fine("# of vertices: " + to_string(Size));
	init_nodes();
	merge_nodes();
	find_clusters();
	compute_quality();
	clogger()->fine("# of clusters: " + to_string(ClusCount));
	clogger()->fine("Quality: " + to_string(Quality));
	if (Quality < DivThreshold && ClusCount > 1)
	{
		ClusCount = 1;
		Labels.assign(Size, 0);
		//Quality = 0;
		clogger()->fine("Quality too low, cancel clustering and merge all clusters");
	}
	if (Subdiv) subdivide();
}

int modularity_clustering::get_cluster_count() const
{
	return ClusCount;
}

const vector<int>& modularity_clustering::get_labels() const
{
	return Labels;
}

double modularity_clustering::get_quality() const
{
	return Quality;
}

void modularity_clustering::init_nodes()
{
	NodeCount = Size;
	Nodes.reserve(Size * 2 - 1);
	Nodes.resize(Size);
	TotalWeight = total_weight();
	for (int I = 0; I < Size; I++)
	{
		Nodes[I].EVal = TotalWeight != 0 ? Weights[I] / TotalWeight : 0;
		Nodes[I].AVal = TotalWeight != 0 ? total_weight(I) / TotalWeight : 0;
		Nodes[I].A = Nodes[I].B = -1;
	}
	Merges.clear();
	Merges.set_size(Size * 2 - 1);
	Heap.clear();
	for (int A = 0; A < Size; A++) for (int I = 0; I < (int)Graph[A].size(); I++)
		if (A < Graph[A][I].to)
	{
		int B = Graph[A][I].to;
		double EVal = Graph[A][I].value / TotalWeight;
		merge Merge;
		Merge.Delta = 2 * (EVal - Nodes[A].AVal * Nodes[B].AVal);
		Merge.A = A;
		Merge.B = B;
		merge_heap::node* HeapNode = Heap.insert(Merge);
		Merges.add_edge(A, B, make_pair((weight_type)EVal, HeapNode));
	}
}

void modularity_clustering::merge_nodes()
{
	vector<char> Marks(Size * 2 - 1, false);
	vector<double> EVals(Size * 2 - 1);
	vector<merge_heap::node*> HeapNodes(Size * 2 - 1);
	vector<int> List;
	List.reserve(Size * 2 - 1);
	while (Heap.size() > 0)
	{
		int A = Heap.top()->key.A;
		int B = Heap.top()->key.B;
		const merge_graph::edge_list &AList = Merges[A], &BList = Merges[B];
		List.clear();
		for (int I = 0; I < (int)AList.size(); I++)
		{
			int C = AList[I].to;
			Marks[C] = true;
			EVals[C] = AList[I].value.first;
			HeapNodes[C] = AList[I].value.second;
			List.push_back(C);
		}
		for (int I = 0; I < (int)BList.size(); I++)
		{
			int C = BList[I].to;
			if (Marks[C] || C == A)
			{
				EVals[C] += BList[I].value.first;
				Heap.erase(BList[I].value.second);
				continue;
			}
			EVals[C] = BList[I].value.first;
			HeapNodes[C] = BList[I].value.second;
			List.push_back(C);
		}
		assert(Marks[B]);
		node Node;
		Node.EVal = EVals[B] * 2 + Nodes[A].EVal + Nodes[B].EVal;
		Node.AVal = Nodes[A].AVal + Nodes[B].AVal;
		Node.A = A;
		Node.B = B;
		Nodes.push_back(Node);
		int AB = NodeCount++;
		Merges.remove_vertex(A);
		Merges.remove_vertex(B);
		Merges[AB].reserve(List.size());
		for (int I = 0; I < (int)List.size(); I++)
		{
			int C = List[I];
			Marks[C] = false;
			if (C == A || C == B) continue;
			merge_heap::node* HeapNode = HeapNodes[C];
			HeapNode->key.Delta = 2 * (EVals[C] - Node.AVal * Nodes[C].AVal);
			HeapNode->key.A = AB;
			HeapNode->key.B = C;
			Heap.update(HeapNode);
			Merges.add_edge(AB, C, make_pair((weight_type)EVals[C], HeapNode));
		}
	}
	assert(NodeCount < Size * 2);
}

void modularity_clustering::find_clusters()
{
	vector<double> Qualities(NodeCount);
	vector<bool> Marks(NodeCount);
	for (int I = 0; I < NodeCount; I++)
	{
		Qualities[I] = Nodes[I].EVal - Nodes[I].AVal * Nodes[I].AVal;
		Marks[I] = true;
		if (I >= Size && Qualities[Nodes[I].A] + Qualities[Nodes[I].B] > Qualities[I])
		{
			Qualities[I] = Qualities[Nodes[I].A] + Qualities[Nodes[I].B];
			Marks[I] = false;
		}
	}
	vector<int> Roots(NodeCount, -1);
	ClusCount = 0;
	for (int I = NodeCount - 1; I >= 0; I--)
	{
		if (Roots[I] == -1 && Marks[I]) Roots[I] = ClusCount++;
		if (I >= Size && Roots[I] >= 0) Roots[Nodes[I].A] = Roots[Nodes[I].B] = Roots[I];
	}
	Labels.assign(Roots.begin(), Roots.begin() + Size);
	reset_vector(Nodes);
	Merges.clear();
	Heap.clear();
}

void modularity_clustering::compute_quality()
{
	Quality = 0;
	if (TotalWeight == 0) return;
	vector<double> EVals(ClusCount, 0), AVals(ClusCount, 0);
	for (int A = 0; A < Size; A++)
	{
		EVals[Labels[A]] += Weights[A];
		AVals[Labels[A]] += Weights[A];
		for (int I = 0; I < (int)Graph[A].size(); I++)
		{
			int B = Graph[A][I].to;
			double Weight = Graph[A][I].value;
			AVals[Labels[A]] += Weight;
			if (Labels[A] == Labels[B]) EVals[Labels[A]] += Weight;
		}
	}
	for (int I = 0; I < ClusCount; I++)
	{
		EVals[I] /= TotalWeight;
		AVals[I] /= TotalWeight;
		Quality += EVals[I] - AVals[I] * AVals[I];
	}
}

void modularity_clustering::subdivide()
{
	if (ClusCount <= 1) return;
	
	logger::task LTask = clogger()->fine("Subdividing clusters");

	vector<int> Indexes(Size);
	vector<int> ClusSizes(ClusCount, 0);
	for (int I = 0; I < Size; I++) Indexes[I] = ClusSizes[Labels[I]]++;

	vector<smart_ptr<modularity_clustering> > Clusters(ClusCount);
	for (int I = 0; I < ClusCount; I++)
	{
		Clusters[I].reset(new modularity_clustering());
		Clusters[I]->set_size(ClusSizes[I]);
		Clusters[I]->set_division_threshold(DivThreshold);
	}
	for (int A = 0; A < Size; A++) for (int I = 0; I < (int)Graph[A].size(); I++)
		if (A <= Graph[A][I].to && Labels[A] == Labels[Graph[A][I].to])
			Clusters[Labels[A]]->add_edge(Indexes[A], Indexes[Graph[A][I].to], Graph[A][I].value);
	for (int A = 0; A < Size; A++) Clusters[Labels[A]]->add_edge(Indexes[A], Indexes[A], Weights[A]);
	for (int I = 0; I < ClusCount; I++) Clusters[I]->compute(true);

	vector<int> Bases(ClusCount + 1);
	Bases[0] = 0;
	for (int I = 1; I <= ClusCount; I++) Bases[I] = Bases[I - 1] + Clusters[I - 1]->get_cluster_count();
	for (int I = 0; I < Size; I++)
		Labels[I] = Bases[Labels[I]] + Clusters[Labels[I]]->Labels[Indexes[I]];
	ClusCount = Bases[ClusCount];
}

double modularity_clustering::total_weight(int A) const
{
	double Result = Weights[A];
	for (int I = 0; I < (int)Graph[A].size(); I++)
		Result += Graph[A][I].value;
	return Result;
}

double modularity_clustering::total_weight() const
{
	double Result = 0;
	for (int I = 0; I < Size; I++) Result += total_weight(I);
	return Result;
}

}

}
